import numpy as np

def read_tum_file(file_path):
    """读取TUM格式轨迹文件：时间戳 x y z qx qy qz qw"""
    data = []
    with open(file_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):  # 跳过空行和注释
                continue
            parts = list(map(float, line.split()))
            if len(parts) != 8:  # TUM格式必须是8列（时间戳+位置+四元数）
                continue
            ts = parts[0]          # 时间戳
            pos = np.array(parts[1:4])  # 位置 (x, y, z)
            quat = np.array(parts[4:8]) # 四元数 (qx, qy, qz, qw)
            data.append((ts, pos, quat))
    return data

def quaternion_linear_interp(q0, q1, alpha):
    """四元数线性插值+归一化（兼容旧scipy版本）
    q0: 起始四元数 (qx, qy, qz, qw)
    q1: 结束四元数 (qx, qy, qz, qw)
    alpha: 插值权重（0~1）
    返回：插值后的归一化四元数
    """
    # 线性插值
    q_interp = (1 - alpha) * q0 + alpha * q1
    # 归一化（确保是合法的旋转四元数，模长=1）
    q_norm = np.linalg.norm(q_interp)
    if q_norm < 1e-6:  # 避免除以0
        return q0
    return q_interp / q_norm

def align_trajectory(truth_data, est_data):
    """将估计轨迹（NDT）插值到真值轨迹的时间戳上，保留所有数据"""
    aligned_est = []
    # 提取真值时间戳、估计轨迹的时间戳/位置/四元数
    truth_ts_list = np.array([d[0] for d in truth_data])
    est_ts_list = np.array([d[0] for d in est_data])
    est_pos_list = np.array([d[1] for d in est_data])
    est_quat_list = np.array([d[2] for d in est_data])
    
    # 对每个真值时间戳，插值估计轨迹的位姿
    for ts in truth_ts_list:
        # 找到估计轨迹中与当前真值时间戳相邻的两个点
        idx = np.searchsorted(est_ts_list, ts)  # 二分查找插入位置
        
        if idx == 0:
            # 真值时间戳早于所有估计时间戳，用第一个估计值
            pos_interp = est_pos_list[0]
            quat_interp = est_quat_list[0]
        elif idx >= len(est_ts_list):
            # 真值时间戳晚于所有估计时间戳，用最后一个估计值
            pos_interp = est_pos_list[-1]
            quat_interp = est_quat_list[-1]
        else:
            # 时间戳在两个估计点之间，线性插值
            t_prev = est_ts_list[idx-1]  # 前一个估计点时间戳
            t_next = est_ts_list[idx]    # 后一个估计点时间戳
            alpha = (ts - t_prev) / (t_next - t_prev)  # 插值权重（0~1）
            
            # 1. 位置线性插值
            pos_interp = (1 - alpha) * est_pos_list[idx-1] + alpha * est_pos_list[idx]
            # 2. 四元数线性插值+归一化
            quat_interp = quaternion_linear_interp(est_quat_list[idx-1], est_quat_list[idx], alpha)
        
        # 格式化为TUM格式（时间戳保留9位小数，其他保留9位）
        line = f"{ts:.9f} {pos_interp[0]:.9f} {pos_interp[1]:.9f} {pos_interp[2]:.9f} " \
               f"{quat_interp[0]:.9f} {quat_interp[1]:.9f} {quat_interp[2]:.9f} {quat_interp[3]:.9f}\n"
        aligned_est.append(line)
    
    return aligned_est

# 主函数：执行对齐
if __name__ == "__main__":
    # --------------------------
    # 请确认你的文件名是否正确！
    # --------------------------
    TRUTH_FILE = "gnss_ground_truth_ndt.txt"  # 真值轨迹文件
    EST_FILE = "laser_odom_ndt.txt"          # NDT估计轨迹文件
    OUTPUT_FILE = "laser_odom_aligned.txt"    # 输出对齐后的NDT轨迹文件
    
    # 1. 读取轨迹数据
    print(f"正在读取文件：{TRUTH_FILE} 和 {EST_FILE}")
    truth_data = read_tum_file(TRUTH_FILE)
    est_data = read_tum_file(EST_FILE)
    
    if len(truth_data) == 0 or len(est_data) == 0:
        print("错误：未读取到轨迹数据，请检查文件路径和格式！")
        exit(1)
    
    # 2. 对齐轨迹（NDT插值到真值时间戳）
    print(f"正在对齐轨迹（共{len(truth_data)}个真值时间戳）...")
    aligned_est_lines = align_trajectory(truth_data, est_data)
    
    # 3. 保存对齐后的文件
    with open(OUTPUT_FILE, 'w') as f:
        f.writelines(aligned_est_lines)
    
    print(f"对齐完成！生成对齐后的NDT轨迹文件：{OUTPUT_FILE}")
    print(f"对齐后文件包含 {len(aligned_est_lines)} 个数据点（与真值数量一致）")
